package controller

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"sugar/database"
	"sugar/errorcode"
	"sugar/response"
	"sugar/utils"

	"gopkg.in/guregu/null.v3"
)

type Budget struct {
	Memory string      `json:"memory"`
	Cpu    string      `json:"cpu"`
	Gpu    null.String `json:"gpu"`
}

//
func (op *Operation) CreateTrainJob(userID string, name string, version int, trainjobType int, projectID null.Int, description null.String, budget null.String, trainDatasetID int64,
	valDatasetID int64, modelID int64, modelCommitID null.String, modelArgs null.String, trainArgs null.String, destScore string, maxTrials int, useAutoml int) *response.Response {
	op.BeginOp(UseK8sOption)
	defer op.EndOp()
	if exist, _ := op.DbOp.IsTrainJobExist(userID, name); exist {
		panic(errorcode.TrainJobAlreadyExistError)
	}
	trainDataset, _ := op.checkDatasetAvailable(userID, trainDatasetID)
	valDataset, _ := op.checkDatasetAvailable(userID, valDatasetID)
	model, _ := op.checkModelAvailable(userID, modelID)
	if !(trainjobType == valDataset.Type && trainjobType == model.Type && trainDataset.Type == trainjobType) {
		panic(errorcode.TypeNotMatchError)
	}
	paramID := utils.GetUUID()
	paramSavePath := "/home/luqin/modelframetest/params/" + paramID
	err := os.Mkdir(paramSavePath, os.ModePerm)
	if err != nil {
		panic(errorcode.MkdirError)
	}
	budgetS := Budget{
		Memory: "1000Mi",
		Cpu:    "2",
	}
	if budget.Valid {
		err = json.Unmarshal([]byte(budget.String), &budgetS)
		if err != nil {
			panic(errorcode.JSONOperationError)
		}
	}
	bbudget, _ := json.Marshal(budgetS)

	modelFolderID := utils.GetUUID()
	err = utils.WriteFile("/home/luqin/modelframetest/model/"+modelFolderID+"/model.py", model.Model_File_Bytes, 0755)
	if err != nil {
		panic(errorcode.FileOperationError)
	}
	// datasetFolderPath:=utils.GetEnvDefault("DATASET_FOLDER_PATH", "/root/dataset")
	trainID, err := op.DbOp.CreateTrainJob(userID, name, version, projectID, description, string(bbudget), trainDatasetID, valDatasetID, modelID, modelCommitID, modelArgs, trainArgs, paramSavePath, useAutoml, trainjobType, destScore, maxTrials)
	if err != nil {
		panic(errorcode.CreateTrainJobError)
	}
	statusHookURL := fmt.Sprint("http://", utils.GetEnvDefault("KUBERNETES_ADVERTISE_ADDR", "192.168.2.89"), ":", utils.GetEnvDefault("KUBERNETES_ADVERTISE_PORT", "8082"), "/", "trainjob/", trainID, "/status")
	err = op.K8sOp.CreateTrainjob(strconv.FormatInt(trainID, 10), paramID, "model", model.Model_Class, modelArgs.String, trainArgs.String, trainDataset.Object, valDataset.Object, destScore, maxTrials, useAutoml, budgetS.Cpu, budgetS.Memory, 0, modelFolderID, statusHookURL)
	if err != nil {
		op.DbOp.UpdateTrainJobStatus(trainID, -1)
		// panic(errorcode.KubernetesOperationError)
	} else {
		op.DbOp.UpdateTrainJobStatus(trainID, 1)
		// go trainJobMonitor(trainID)
	}

	res := make(map[string]interface{})
	res["id"] = trainID
	return response.NewSuccess(res)
}

func trainJobMonitor(trainJobID int64) {
	op := NewOp()
	op.updateTrainJob(trainJobID)
}

func (op *Operation) updateTrainJob(trainJobID int64) {
	sTrainJobID := strconv.FormatInt(trainJobID, 10)
	var status string
	for {
		op.BeginOp(UseK8sOption)
		status = op.K8sOp.GetTrainjobStatus(sTrainJobID)
		if status == "TERMINATED" {
			_ = op.DbOp.UpdateTrainJobStatus(trainJobID, 3)
		} else if status == "FAILED" {
			_ = op.DbOp.UpdateTrainJobStatus(trainJobID, -1)
		} else if status == "SUCCESS" {
			_ = op.DbOp.UpdateTrainJobStatus(trainJobID, 2)
		} else {
			_ = op.DbOp.UpdateTrainJobStatus(trainJobID, 1)
		}
		op.EndOp()
		if status == "TERMINATED" || status == "FAILED" || status == "SUCCESS" {
			break
		}
	}
}

func (op *Operation) checkTrainJobAuth(userID string, trainJobID int64) (*database.Trainjob, error) {
	trainJob, err := op.DbOp.GetTrainjobByID(trainJobID)
	if err != nil {
		panic(errorcode.GetTrainJobError)
	}
	if trainJob.User_ID.String() != userID {
		panic(errorcode.NoTrainJobAuthError)
	}
	return trainJob, nil
}

//
func (op *Operation) GetTrainJob(userID string, trainJobID int64) *response.Response {
	op.BeginOp()
	defer op.EndOp()
	trainJob, _ := op.checkTrainJobAuth(userID, trainJobID)
	return response.NewSuccess(trainJob)
}

//
func (op *Operation) StopTrainJob(userID string, trainJobID int64) *response.Response {
	op.BeginOp(UseK8sOption)
	defer op.EndOp()
	op.checkTrainJobAuth(userID, trainJobID)
	err := op.K8sOp.DeleteTrainjob(strconv.FormatInt(trainJobID, 10))
	if err != nil {
		panic(errorcode.DeleteTrainJobError)
	}
	return response.NewSuccess(nil)
}

//
func (op *Operation) ListTrainJob(userID string) *response.Response {
	op.BeginOp()
	defer op.EndOp()
	trainJobs, err := op.DbOp.ListTrainjobByUser(userID)
	if err != nil {
		panic(errorcode.ListTrainJobError)
	}
	return response.NewSuccess(trainJobs)
}

//
func (op *Operation) UpdateTrainJobStatus(trainJobID int64, status string) *response.Response {
	op.BeginOp()
	defer op.EndOp()
	var istatus int
	switch status {
	case "TERMINATED":
		istatus = 3
		break
	case "FAILED":
		istatus = -1
		break
	case "SUCCESS":
		istatus = 2
		break
	default:
		istatus = 1
		break
	}
	if err := op.DbOp.UpdateTrainJobStatus(trainJobID, istatus); err != nil {
		panic(errorcode.UpdateTrainJobStatusError)
	}
	return response.NewSuccess(nil)
}
